// prefetcher.cc (Version 4.6 - The Final Polish)
#include <iostream>
#include <fstream>
#include <string>
#include <vector>
#include <sstream>
#include <csignal>
#include <ctime>
#include <unistd.h>
#include <sys/syscall.h>
#include <fcntl.h>
#include <cerrno>
#include <cstring>
#include <unordered_map>
#include <iomanip>
#include <deque>
#include <utility>
#include <algorithm>
#include <cmath>
#include <random>
#include <chrono>

extern "C" {
#include <bpf/libbpf.h>
#include <bpf/bpf.h>
#include "prefetcher.skel.h"
#include "prefetcher.h"
}

// ============================================================================
// V4.6: 模型核心
// ============================================================================
#define HISTORY_SIZE 5
#define ANALYSIS_THRESHOLD 3
#define CONFIDENCE_THRESHOLD 3      // V4.6: 恢复到更稳健的阈值
#define RATE_TOLERANCE 0.30         // V4.6: 容忍度放宽到30%
#define INITIAL_PREFETCH_HORIZON_S 0.8
#define MAX_PREFETCH_HORIZON_S 1.5
#define MIN_PREFETCH_HORIZON_S 0.1
#define EWMA_ALPHA 0.3
#define MEM_PRESSURE_THRESHOLD_KB (512 * 1024)
#define STABLE_PROBE_THRESHOLD 5
#define PREFETCH_COOLDOWN_FAULTS 2  // V4.6: 预取后忽略2次缺页

struct FaultRecord {
    struct timespec timestamp;
    off_t file_offset;
};

struct StreamTracker {
    uint32_t pid;
    uint32_t tid;
    std::deque<FaultRecord> history;
    std::string file_path;
    int direction;
    double ewma_rate_bps;
    int confidence;
    double current_horizon_s;
    struct timespec last_prefetch_ts;
    bool awaiting_feedback;
    int stable_feedback_count;
    int cooldown_faults_remaining; // V4.6: 新增冷却计数器
    StreamTracker(uint32_t p, uint32_t t, const std::string& path)
        : pid(p), tid(t), file_path(path), direction(0), ewma_rate_bps(0.0),
          confidence(0), current_horizon_s(INITIAL_PREFETCH_HORIZON_S),
          awaiting_feedback(false), stable_feedback_count(0), cooldown_faults_remaining(0) {
        last_prefetch_ts.tv_sec = 0;
        last_prefetch_ts.tv_nsec = 0;
    }
};

static std::unordered_map<uint32_t, StreamTracker> stream_trackers;
// ============================================================================
// 全局变量和辅助函数
// ============================================================================
std::ofstream log_file;
static volatile bool exiting = false;
static long faults_processed_count = 0;
static long faults_in_files_count = 0;
static long pattern_detected_count = 0;
static long readahead_success_count = 0;
static long readahead_failure_count = 0;
static std::mt19937 rng(std::chrono::steady_clock::now().time_since_epoch().count());
static std::uniform_real_distribution<> distrib(0.0, 1.0);

double get_elapsed_seconds(const struct timespec& start, const struct timespec& end) {
    return (end.tv_sec - start.tv_sec) + (end.tv_nsec - start.tv_nsec) / 1e9;
}
std::string get_timestamp() {
    time_t now = time(0);
    struct tm tstruct;
    char buf[80];
    tstruct = *localtime(&now);
    strftime(buf, sizeof(buf), "%Y-%m-%d %X", &tstruct);
    return buf;
}
void log_message(const std::string& msg) {
    std::string full_msg = "[" + get_timestamp() + "] " + msg;
    std::cout << full_msg << std::endl;
    if (log_file.is_open()) {
        log_file << full_msg << std::endl;
    }
}
long get_available_memory_kb() {
    std::ifstream meminfo("/proc/meminfo");
    std::string line;
    long mem_available = -1;
    while (std::getline(meminfo, line)) {
        if (line.rfind("MemAvailable:", 0) == 0) {
            std::stringstream ss(line);
            std::string key;
            ss >> key >> mem_available;
            break;
        }
    }
    return mem_available;
}
bool find_file_offset(uint32_t pid, unsigned long vma_start, int& fd, off_t& file_offset, std::string& file_path) {
    std::string maps_path = "/proc/" + std::to_string(pid) + "/maps";
    std::ifstream maps_file(maps_path);
    if (!maps_file.is_open()) return false;
    std::string line;
    while (std::getline(maps_file, line)) {
        std::stringstream ss(line);
        std::string address_range, perms, offset_str, dev, inode, pathname;
        ss >> address_range >> perms >> offset_str >> dev >> inode >> pathname;
        if (pathname.empty() || pathname[0] == '[') continue;
        
        unsigned long start, end, file_map_offset;
        sscanf(address_range.c_str(), "%lx-%lx", &start, &end);
        sscanf(offset_str.c_str(), "%lx", &file_map_offset);

        if (vma_start >= start && vma_start < end) {
            fd = open(pathname.c_str(), O_RDONLY);
            if (fd < 0) return false;
            file_offset = (vma_start - start) + file_map_offset;
            file_path = pathname;
            return true;
        }
    }
    return false;
}
// ============================================================================
// V4.6: 核心事件处理逻辑
// ============================================================================
int handle_event(void *ctx, void *data, size_t data_sz) {
    const struct fault_info *e = (struct fault_info*)data;
    faults_processed_count++;
    struct timespec ts_now;
    clock_gettime(CLOCK_MONOTONIC, &ts_now);
    int fd = -1;
    off_t current_offset;
    std::string file_path;
    if (!find_file_offset(e->pid, e->addr, fd, current_offset, file_path)) {
        if (fd != -1) close(fd);
        return 0;
    }
    faults_in_files_count++;
    auto it = stream_trackers.find(e->tid);
    if (it == stream_trackers.end() || it->second.file_path != file_path) {
        if (it != stream_trackers.end()) stream_trackers.erase(it);
        it = stream_trackers.emplace(e->tid, StreamTracker(e->pid, e->tid, file_path)).first;
        log_message("New stream for TID " + std::to_string(e->tid) + " on file: " + file_path);
    }
    
    StreamTracker& tracker = it->second;

    // V4.6: Cooldown Logic
    if (tracker.cooldown_faults_remaining > 0) {
        tracker.cooldown_faults_remaining--;
        log_message("TID " + std::to_string(e->tid) + ": In cooldown (" + std::to_string(tracker.cooldown_faults_remaining) + " faults left). Skipping analysis.");
        if (fd != -1) close(fd);
        return 0;
    }

    if (tracker.awaiting_feedback) {
        double elapsed = get_elapsed_seconds(tracker.last_prefetch_ts, ts_now);
        std::stringstream log_msg;
        log_msg << "TID " << e->tid << ": Feedback. Elapsed: " << std::fixed << std::setprecision(3) << elapsed << "s. ";
        bool adjusted = false;
        if (elapsed > tracker.current_horizon_s * 2.0) {
            tracker.current_horizon_s = std::max(MIN_PREFETCH_HORIZON_S, tracker.current_horizon_s * 0.9);
            adjusted = true;
        } else if (elapsed < tracker.current_horizon_s * 0.5) {
            tracker.current_horizon_s = std::min(MAX_PREFETCH_HORIZON_S, tracker.current_horizon_s * 1.1);
            adjusted = true;
        }

        if(adjusted) {
            log_msg << "Horizon -> " << tracker.current_horizon_s << "s.";
            tracker.stable_feedback_count = 0;
        } else {
            tracker.stable_feedback_count++;
            log_msg << "Horizon stable (" << tracker.stable_feedback_count << ").";
        }
        log_message(log_msg.str());
        
        if(tracker.stable_feedback_count >= STABLE_PROBE_THRESHOLD) {
           double probe_factor = (distrib(rng) > 0.5) ? 1.20 : 0.80; // V4.6: More aggressive probe
           tracker.current_horizon_s *= probe_factor;
           tracker.current_horizon_s = std::max(MIN_PREFETCH_HORIZON_S, std::min(MAX_PREFETCH_HORIZON_S, tracker.current_horizon_s));
           log_message("  -> Probing new horizon: " + std::to_string(tracker.current_horizon_s) + "s.");
           tracker.stable_feedback_count = 0;
        }
        tracker.awaiting_feedback = false;
    }
    
    tracker.history.push_back({ts_now, current_offset});
    if (tracker.history.size() > HISTORY_SIZE) tracker.history.pop_front();
    if (tracker.history.size() < ANALYSIS_THRESHOLD) {
        if (fd != -1) close(fd);
        return 0;
    }

    const auto& first_record = tracker.history.front();
    const auto& last_record = tracker.history.back();
    double delta_time = get_elapsed_seconds(first_record.timestamp, last_record.timestamp);
    if (delta_time < 1e-6) {
        if (fd != -1) close(fd);
        return 0;
    }
    off_t delta_offset = last_record.file_offset - first_record.file_offset;
    double current_rate = delta_offset / delta_time;
    int current_direction = (delta_offset > 0) ? 1 : 0;
    
    if (tracker.ewma_rate_bps < 1.0) tracker.ewma_rate_bps = current_rate;
    else tracker.ewma_rate_bps = (current_rate * EWMA_ALPHA) + (tracker.ewma_rate_bps * (1.0 - EWMA_ALPHA));
    
    double rate_diff = std::abs(current_rate - tracker.ewma_rate_bps) / tracker.ewma_rate_bps;

    if (current_direction != tracker.direction) {
        tracker.confidence = 0;
    } else if (rate_diff < RATE_TOLERANCE) {
        tracker.confidence++;
    } else {
        tracker.confidence = std::max(0, tracker.confidence - 1);
    }
    tracker.direction = current_direction;
    
    log_message("TID " + std::to_string(e->tid) + ": Rate(inst/ewma): " + 
                std::to_string((int)(current_rate/(1024*1024))) + "/" + 
                std::to_string((int)(tracker.ewma_rate_bps/(1024*1024))) + " MB/s, " +
                "Confidence: " + std::to_string(tracker.confidence) + ", Direction: " + std::to_string(tracker.direction));

    if (tracker.confidence >= CONFIDENCE_THRESHOLD && tracker.direction == 1) {
        pattern_detected_count++;
        long mem_avail_kb = get_available_memory_kb();
        double load_factor = 1.0;
        if (mem_avail_kb != -1 && mem_avail_kb < MEM_PRESSURE_THRESHOLD_KB) {
            load_factor = (double)mem_avail_kb / MEM_PRESSURE_THRESHOLD_KB;
            log_message("  -> MEMORY PRESSURE! Throttling prefetch with factor: " + std::to_string(load_factor));
        }

        size_t prefetch_bytes = tracker.ewma_rate_bps * tracker.current_horizon_s * load_factor;
        // V4.6: More conservative prefetch cap
        prefetch_bytes = std::max((size_t) (256 * 1024), std::min(prefetch_bytes, (size_t) (8 * 1024 * 1024)));
        off_t prefetch_offset = last_record.file_offset;
        
        long ret = syscall(SYS_readahead, fd, prefetch_offset, prefetch_bytes);
        
        if (ret == 0) {
            readahead_success_count++;
            log_message("  -> STEADY PATTERN! Prefetching " + std::to_string(prefetch_bytes / 1024) + " KB. Horizon: " + std::to_string(tracker.current_horizon_s) + "s.");
            tracker.last_prefetch_ts = ts_now;
            tracker.awaiting_feedback = true;
            tracker.history.clear();
            tracker.confidence = 0;
            tracker.cooldown_faults_remaining = PREFETCH_COOLDOWN_FAULTS; // V4.6: Activate cooldown
        } else {
            readahead_failure_count++;
            log_message("  -> readahead() FAILED. Errno: " + std::to_string(errno) + " (" + strerror(errno) + ")");
        }
    }

    if (fd != -1) close(fd);
    return 0;
}
// ============================================================================
// Main 函数
// ============================================================================
int print_libbpf_log(enum libbpf_print_level level, const char *format, va_list args) {
    if (level == LIBBPF_WARN) {
        char buffer[1024];
        vsnprintf(buffer, sizeof(buffer), format, args);
        log_message("libbpf: " + std::string(buffer));
    }
    return 0;
}
static void sig_handler(int sig) { exiting = true; }
int main(int argc, char **argv) {
    log_file.open("prefetcher.log", std::ios::out | std::ios::trunc);
    log_message("Prefetcher V4.6 (Final Polish) starting...");
    libbpf_set_print(print_libbpf_log);
    struct prefetcher_bpf *skel = prefetcher_bpf__open_and_load();
    if (!skel) { 
        log_message("Error: Failed to open and load BPF skeleton"); 
        return 1; 
    }
    struct config cfg = {};
    if (argc > 1) {
        cfg.target_pid = std::stoi(argv[1]);
        log_message("Targeting PID: " + std::string(argv[1]));
    } else {
        cfg.target_pid = 0;
        log_message("Monitoring all processes.");
    }
    const int key = 0;
    int err = bpf_map_update_elem(bpf_map__fd(skel->maps.config_map), &key, &cfg, BPF_ANY);
    if (err) {
        log_message("Error: Failed to update config map: " + std::string(strerror(-err)));
        prefetcher_bpf__destroy(skel);
        return 1;
    }
    err = prefetcher_bpf__attach(skel);
    if (err) { 
        log_message("Error: Failed to attach BPF skeleton: " + std::string(strerror(-err))); 
        prefetcher_bpf__destroy(skel); 
        return 1; 
    }
    log_message("eBPF program loaded and attached successfully.");
    signal(SIGINT, sig_handler); signal(SIGTERM, sig_handler);
    struct ring_buffer *rb = ring_buffer__new(bpf_map__fd(skel->maps.events), handle_event, NULL, NULL);
    if (!rb) { 
        err = -1; 
        log_message("Error: Failed to create ring buffer"); 
        goto cleanup; 
    }
    log_message("Waiting for events... Press Ctrl+C to exit.");
    while (!exiting) {
        err = ring_buffer__poll(rb, 100);
        if (err == -EINTR) { err = 0; break; }
        if (err < 0) { log_message("Error polling ring buffer: " + std::to_string(err)); break; }
    }
cleanup:
    ring_buffer__free(rb);
    prefetcher_bpf__destroy(skel);
    log_message("Prefetcher stopped. Final statistics:");
    std::stringstream final_stats;
    final_stats << "  - Total page faults processed: " << faults_processed_count << std::endl;
    final_stats << "  - Page faults on file mappings: " << faults_in_files_count << std::endl;
    final_stats << "  - Steady patterns detected: " << pattern_detected_count << std::endl;
    final_stats << "  - Successful readahead calls: " << readahead_success_count << std::endl;
    final_stats << "  - Failed readahead calls: " << readahead_failure_count;
    log_message(final_stats.str());
    if (log_file.is_open()) log_file.close();
    return err < 0 ? -err : 0;
}